# import os
# import pandas as pd
# import numpy as np
# from scipy.optimize import curve_fit
# import matplotlib.pyplot as plt
#
# # 设置中文显示
# plt.rcParams['font.sans-serif'] = ['Microsoft YaHei']  # 微软雅黑
# plt.rcParams['axes.unicode_minus'] = False    # 解决负号乱码
#
# # 定义幂函数模型（H = aQ^b）
# def h_q_model(Q, a, b):
#     return a * Q ** b
#
# # 输入文件夹路径（根据实际修改）
# folder_path = r"F:\研究生\项目\深圳万泉河\马斯京跟参数计算\QH曲线"
#
# # 获取所有Excel文件
# excel_files = [f for f in os.listdir(folder_path) if f.endswith('.xlsx')]
#
# # 创建结果汇总表
# results = []
#
# for file in excel_files:
#     file_path = os.path.join(folder_path, file)
#     try:
#         # 读取数据
#         df = pd.read_excel(file_path)
#         Q = df['Q (m³/s)'].values  # Q作为自变量（x轴）
#         H = df['H (m)'].values    # H作为因变量（y轴）
#
#         # 数据清洗：去除零值、负值和空值
#         valid_mask = (Q > 0) & (H > 0) & (~np.isnan(Q)) & (~np.isnan(H))
#         Q_clean = Q[valid_mask]
#         H_clean = H[valid_mask]
#
#         if len(Q_clean) < 3:
#             raise ValueError("有效数据点不足，无法拟合")
#
#         # 尝试幂函数拟合（H = aQ^b）
#         params, covariance = curve_fit(h_q_model, Q_clean, H_clean, maxfev=10000)
#         a, b = params
#         model_type = 'power'
#
#         # 计算拟合优度 R²
#         H_fit = h_q_model(Q_clean, a, b)
#         residuals = H_clean - H_fit
#         ss_res = np.sum(residuals ** 2)
#         ss_tot = np.sum((H_clean - np.mean(H_clean)) ** 2)
#         r_squared = 1 - (ss_res / ss_tot)
#
#     except Exception as e:
#         print(f"幂函数拟合失败，尝试多项式拟合: {file}, 错误: {str(e)}")
#         try:
#             # 多项式拟合（H = aQ² + bQ）
#             coeffs = np.polyfit(Q_clean, H_clean, 2)
#             a_poly, b_poly, c_poly = coeffs
#             a, b = a_poly, b_poly  # 仅保存前两个系数（忽略常数项c）
#             model_type = 'poly'
#
#             # 计算拟合值和 R²
#             H_fit = a_poly * Q_clean**2 + b_poly * Q_clean + c_poly
#             residuals = H_clean - H_fit
#             ss_res = np.sum(residuals ** 2)
#             ss_tot = np.sum((H_clean - np.mean(H_clean)) ** 2)
#             r_squared = 1 - (ss_res / ss_tot)
#
#         except Exception as e:
#             print(f"多项式拟合也失败: {file}, 错误: {str(e)}")
#             continue
#
#     # 保存结果
#     results.append({
#         '文件名': file,
#         '参数a': round(a, 4),
#         '参数b': round(b, 4),
#         'R²': round(r_squared, 4),
#         '模型类型': model_type
#     })
#
#     # 绘图（Q为x轴，H为y轴）
#     plt.figure()
#     plt.scatter(Q_clean, H_clean, label='实测数据', color='blue')
#     if model_type == 'power':
#         plt.plot(Q_clean, h_q_model(Q_clean, a, b), 'r-',
#                  label=f'幂函数拟合: H={a:.2f}Q^{b:.2f}')
#     else:
#         plt.plot(Q_clean, H_fit, 'r-',
#                  label=f'多项式拟合: H={a:.2f}Q² + {b:.2f}Q')
#     plt.xlabel('流量 (m³/s)')
#     plt.ylabel('水位 (m)')
#     plt.legend()
#     plt.title(f'文件 {os.path.splitext(file)[0]} 拟合结果')
#     plt.savefig(os.path.join(folder_path, f'{os.path.splitext(file)[0]}.png'))
#     plt.close()
#
# # 输出结果到Excel
# if results:
#     results_df = pd.DataFrame(results)
#     results_df.to_excel(os.path.join(folder_path, 'QH曲线拟合结果汇总.xlsx'), index=False)
#     print("拟合完成！结果已保存到：QH曲线拟合结果汇总.xlsx")
# else:
#     print("无有效数据完成拟合！")


import os
import pandas as pd
import numpy as np
from scipy.optimize import curve_fit
import matplotlib.pyplot as plt

# 设置中文显示
plt.rcParams['font.sans-serif'] = ['Microsoft YaHei']  # 微软雅黑
plt.rcParams['axes.unicode_minus'] = False    # 解决负号乱码

# 定义幂函数模型（H = aQ^b）
def h_q_model(Q, a, b):
    return a * Q ** b

# 输入文件夹路径（根据实际修改）
folder_path = r"F:\研究生\项目\深圳万泉河\马斯京跟参数计算\QH曲线"

# 获取所有Excel文件
excel_files = [f for f in os.listdir(folder_path) if f.endswith('.xlsx')]

# 创建结果汇总表
results = []

for file in excel_files:
    file_path = os.path.join(folder_path, file)
    try:
        # 读取数据
        df = pd.read_excel(file_path)
        Q = df['Q (m³/s)'].values  # Q作为自变量（x轴）
        H = df['H (m)'].values    # H作为因变量（y轴）

        # 数据清洗：去除零值、负值和空值
        valid_mask = (Q > 0) & (H > 0) & (~np.isnan(Q)) & (~np.isnan(H))
        Q_clean = Q[valid_mask]
        H_clean = H[valid_mask]

        if len(Q_clean) < 3:  # 二次多项式至少需要3个点
            raise ValueError("有效数据点不足，无法拟合二次多项式")

        # 尝试幂函数拟合（H = aQ^b）
        params, covariance = curve_fit(h_q_model, Q_clean, H_clean, maxfev=10000)
        a_power, b_power = params
        model_type = 'power'

        # 计算拟合优度 R²
        H_fit_power = h_q_model(Q_clean, a_power, b_power)
        residuals_power = H_clean - H_fit_power
        ss_res_power = np.sum(residuals_power ** 2)
        ss_tot_power = np.sum((H_clean - np.mean(H_clean)) ** 2)
        r_squared_power = 1 - (ss_res_power / ss_tot_power)

        # 判断R²是否达标
        if r_squared_power < 0.7:
            raise ValueError(f"R²={r_squared_power:.2f} < 0.7，切换二次多项式拟合")

    except Exception as e:
        print(f"幂函数拟合失败或R²不足，尝试二次多项式拟合: {file}, 错误: {str(e)}")
        try:
            # 二次多项式拟合（H = aQ² + bQ + c）
            coeffs = np.polyfit(Q_clean, H_clean, 2)
            a_poly, b_poly, c_poly = coeffs
            model_type = 'poly2'

            # 计算拟合值和 R²
            H_fit_poly = a_poly * Q_clean**2 + b_poly * Q_clean + c_poly
            residuals_poly = H_clean - H_fit_poly
            ss_res_poly = np.sum(residuals_poly ** 2)
            ss_tot_poly = np.sum((H_clean - np.mean(H_clean)) ** 2)
            r_squared_poly = 1 - (ss_res_poly / ss_tot_poly)

        except Exception as e:
            print(f"二次多项式拟合也失败: {file}, 错误: {str(e)}")
            continue

    # 保存结果
    results.append({
        '文件名': file,
        '参数a': round(a_poly if model_type == 'poly2' else a_power, 4),
        '参数b': round(b_poly if model_type == 'poly2' else b_power, 4),
        '参数c': round(c_poly if model_type == 'poly2' else 0, 4),
        'R²': round(r_squared_poly if model_type == 'poly2' else r_squared_power, 4),
        '模型类型': model_type
    })

    # 绘图（Q为x轴，H为y轴）
    plt.figure()
    plt.scatter(Q_clean, H_clean, label='实测数据', color='blue')
    if model_type == 'power':
        plt.plot(Q_clean, h_q_model(Q_clean, a_power, b_power), 'r-',
                 label=f'幂函数拟合: H={a_power:.2f}Q^{b_power:.2f}')
    else:
        plt.plot(Q_clean, a_poly * Q_clean**2 + b_poly * Q_clean + c_poly, 'r-',
                 label=f'二次多项式拟合: H={a_poly:.2f}Q² + {b_poly:.2f}Q + {c_poly:.2f}')
    plt.xlabel('流量 (m³/s)')
    plt.ylabel('水位 (m)')
    plt.legend()
    plt.title(f'文件 {os.path.splitext(file)[0]} 拟合结果')
    plt.savefig(os.path.join(folder_path, f'{os.path.splitext(file)[0]}.png'))
    plt.close()

# 输出结果到Excel
if results:
    results_df = pd.DataFrame(results)
    results_df.to_excel(os.path.join(folder_path, 'QH曲线拟合结果汇总.xlsx'), index=False)
    print("拟合完成！结果已保存到：QH曲线拟合结果汇总.xlsx")
else:
    print("无有效数据完成拟合！")